from typing import List, Optional

# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    @staticmethod
    def first_match(lst, func):
        for i, e in enumerate(lst):
            if func(e):
                return i

    def buildTree(self, preorder: List[int], inorder: List[int]) -> Optional[TreeNode]:
        if len(preorder) == 0:
            return None
        self.preorder = preorder
        self.inorder = inorder
        i_index = inorder.index(preorder[0])
        return self._build_tree(0, i_index, left_i=0, right_i=len(preorder))
    
    def _build_tree(self, p_index, i_index,  left_i, right_i):
        root = TreeNode(self.preorder[p_index])
        li, ri = self.inorder[left_i:i_index], self.inorder[i_index+1:right_i]
        pp = self.preorder[p_index+1:]
        # print('LR: ', li, ri)
        if len(li) != 0:
            next_p = p_index + 1 + self.first_match(pp, lambda x: x in li)
            next_i = left_i + li.index(self.preorder[next_p])
            root.left = self._build_tree(next_p, next_i, left_i=left_i, right_i=i_index)
        if len(ri) != 0:
            next_p = p_index + 1 + self.first_match(pp, lambda x: x in ri)
            next_i = i_index + 1 + ri.index(self.preorder[next_p])
            root.right = self._build_tree(next_p, next_i, left_i=i_index+1, right_i=right_i)
        return root

if __name__ == '__main__':
    s = Solution()
    s.buildTree(preorder = [3,9,20,15,7], inorder = [9,3,15,20,7])
